import cv2
import numpy as np
import os
import sys
import re


def read_calib(path):
    calib_params = {}
    cv_file = cv2.FileStorage(path, cv2.FILE_STORAGE_READ) 
    calib_params = {
        "mtxL": cv_file.getNode("mtxL").mat(), # Camera Matrix Left Pre-Rectification
        "distL": cv_file.getNode("distL").mat(), # Distorsion Left Pre-Rectification
        "rectL": cv_file.getNode("rectL").mat(), # Rectification Transformation for Left Camera
        "proj_matL": cv_file.getNode("proj_matL").mat(), # Camera Matrix Post-Rectification (Use this)
        "wL": int(cv_file.getNode("wL").real()), # Left Image Width
        "hL": int(cv_file.getNode("hL").real()), # Left Image Height

        "mtxR": cv_file.getNode("mtxR").mat(), # Camera Matrix Right Pre-Rectification
        "distR": cv_file.getNode("distR").mat(), # Distorsion Right Pre-Rectification
        "rectR": cv_file.getNode("rectR").mat(), # Rectification Right Transformation for Right Camera
        "proj_matR": cv_file.getNode("proj_matR").mat(), # Camera Matrix Right Post-Rectification (Use this)
        "wR": int(cv_file.getNode("wR").real()), # Right Image Width
        "hR": int(cv_file.getNode("hR").real()), # Right Image Height

        "f_disp": cv_file.getNode("proj_matL").mat()[0,0], # Focal Lenght (Post Rectification)
        "cx": cv_file.getNode("proj_matL").mat()[0,2], # CX (Post Rectification)
        "cy": cv_file.getNode("proj_matL").mat()[1,2], # CY (Post Rectification)
        "b_disp": float(cv_file.getNode("baselineLR").real()) / 1000. # Baseline Left-Right in meters
    }
    cv_file.release()
    return calib_params
    
def invert_d(d, f, b):
    """Convert disparity to depth and viceversa
    Args:
        d: disparity or depth
        f: focal lenght in pixel
        b: baseline
    """
    d = d.astype(np.float32)
    d[d<0] = 0
    d = f * b / d
    return d

def decode_3_channels(raw, max_depth=1000):
    """Carla format to depth
    Args:
        raw: carla format depth image. Expected in BGR.
        max_depth: max depth used during rendering
    """
    raw = raw.astype(np.float32)
    out = raw[:,:,2] + raw[:,:,1] * 256 + raw[:,:,0]*256*256
    out = out / (256*256*256 - 1) * max_depth
    return out


def readPFM(file):
    file = open(file, "rb")
    header = file.readline().rstrip()

    if (sys.version[0]) == "3":
        header = header.decode("utf-8")
    if header == "PF":
        color = True
    elif header == "Pf":
        color = False
    else:
        raise Exception("Not a PFM file.")

    if (sys.version[0]) == "3":
        dim_match = re.match(r"^(\d+)\s(\d+)\s$", file.readline().decode("utf-8"))
    else:
        dim_match = re.match(r"^(\d+)\s(\d+)\s$", file.readline())
    if dim_match:
        width, height = map(int, dim_match.groups())
    else:
        raise Exception("Malformed PFM header.")

    if (sys.version[0]) == "3":
        scale = float(file.readline().rstrip().decode("utf-8"))
    else:
        scale = float(file.readline().rstrip())

    if scale < 0:  # little-endian
        endian = "<"
        scale = -scale
    else:
        endian = ">"  # big-endian

    data = np.fromfile(file, endian + "f")
    shape = (height, width, 3) if color else (height, width)

    data = np.reshape(data, shape)
    data = np.flipud(data)
    return data


def read_d(path, invert=False, f=1, b=1, scale_factor=32):
    """Read depth or disp Map
    Args:
        path: path to depth or disp
        invert: convert d into darity and viceversa
        f_d: focal to convert darity to d
        b_d: baseline to convert darity to d
        scale_factor: scale factor used to decode png 16 bit images
    """

    if path.endswith("pfm"):
        d = readPFM(path)
    elif path.endswith("npy"):
        d = np.load(path)
    elif path.endswith("exr"):
        d = cv2.imread(path, cv2.IMREAD_UNCHANGED)
        d = d[:,:,0]
    elif path.endswith("png"):
        d = cv2.imread(path, cv2.IMREAD_UNCHANGED)
        if len(d.shape) == 3:
            d = decode_3_channels(d)
        elif d.dtype == np.uint16:
            d = d.astype(np.float32)
            d = d / scale_factor
    else:
        d = cv2.imread(path)[:,:,0]
    if invert:
        d = invert_d(d, f, b)
    return d


def save_d(path, d, scale_factor=32):
    """Save depth or disp Map
    Args:
        path: path to output
        d: disparity or depth
        scale_factor: scale factor used for png 16 bit
    """
    _, formt = os.path.splitext(path)
    if formt == ".npy":
        np.save(path, d)
    elif formt == ".png":
        d = d.astype(float)
        d = d * scale_factor
        d = d.astype(np.uint16)
        cv2.imwrite(path, d)
    else:
        cv2.imwrite(path, d)